# -*- coding: utf-8 -*-
# @Time:    2023-05-11 14:45
# @Author:  geng
# @Email:   yonglonggeng@163.com
# @WeChat:  superior_god
# @File:    main.py
# @Project: seg_data_process
# @Package: 
# @Ref:
import json
import os
import xml.etree.ElementTree as ET

labels = set()
template = {
    "version": "4.0.0",
    "flags": {},
    "shapes": [
        {
            "label": "person",
            "points": [
                [
                    191.0,
                    107.36900369003689
                ],
                [
                    313.0,
                    329.36900369003695
                ]
            ],
            "group_id": None,
            "shape_type": "rectangle",
            "flags": {}
        }
    ],
    "imagePath": "2011_000003.jpg",
    "imageData": None,
    "imageHeight": 338,
    "imageWidth": 500
}


def convert_single_xml2json(src):
    xml_path = src
    tree = ET.parse(xml_path)
    filename = tree.findtext("filename")
    size = tree.find("size")

    template['imagePath'] = filename.split(".")[0] + ".jpg"
    template['imageHeight'] = int(size.findtext("height"))
    template['imageWidth'] = int(size.findtext("width"))

    shapes = []
    shape = {
        "label": "person",
        "points": [
            [
                191.0,
                107.36900369003689
            ],
            [
                313.0,
                329.36900369003695
            ]
        ],
        "group_id": None,
        "shape_type": "rectangle",
        "flags": {}
    }
    all_object = tree.findall("object")
    for obj in all_object:
        shape['label'] = obj.findtext("name")
        labels.add(shape['label'])
        shape['group_id'] = None
        shape['shape_type'] = "rectangle"
        shape['flags'] = {}

        bndbox = obj.find("bndbox")
        shape['points'] = [[int(bndbox.findtext("xmin")), int(bndbox.findtext("ymin"))],
                           [int(bndbox.findtext("xmax")), int(bndbox.findtext("ymax"))]]
        shapes.append(shape)

    template['shapes'] = shapes
    # print(template)

    json.dump(template, open((src.split("/")[-1]).split(".")[0] + ".json", 'w'), indent=2)


if __name__ == '__main__':
    src = "xml/01_bpmh_13.xml"
    des = "01_bpmh_1.data_annotated"
    for file in os.listdir("xml"):
        convert_single_xml2json("xml/" + file)
    print("labels:", labels)
